//
// Created by alax on 4/13/17.
//
#include "testFunctions.h"
#include "llvm/CodeGen/MachineInstrBuilder.h"
#include "llvm/MC/MCContext.h"
#include "llvm/CodeGen/MachineModuleInfo.h"
#include "X86MachineFunctionInfo.h"
using namespace llvm;

namespace {
    class InsertVirtMMUInstrPass:public MachineFunctionPass{
    public:
        static char ID;
        InsertVirtMMUInstrPass():MachineFunctionPass(ID){};
        bool runOnMachineFunction(MachineFunction &MF) override;
        bool handleMainFunction(MachineFunction &MF);
        bool handleNormalFunction(MachineFunction &MF);
        bool processInstruction(MachineBasicBlock &MBB,MachineInstr &MI);

        //processBasicBlock
        bool handlePrologue(MachineFunction::iterator MFI);
        bool handleBlock(MachineFunction::iterator MFI);
        bool handleEPrologue(MachineFunction::iterator MFI);

        //processInstruction
        bool handleInstruction(MachineBasicBlock &MBB,MachineBasicBlock::iterator MBI);
        bool insertDereferenceInstr(MachineBasicBlock &MBB,MachineBasicBlock::iterator MBI);
        bool handleControlFlowInstruction(MachineBasicBlock &MBB,MachineBasicBlock::iterator MBI);

        //processStack
        bool alignStack(MachineBasicBlock &MBB,MachineBasicBlock::iterator MBI);
        bool epalignStack(MachineBasicBlock &MBB,MachineBasicBlock::iterator MBI);

    private:
        const TargetInstrInfo *TII;
        const TargetRegisterInfo *TRI;
    };
}


bool InsertVirtMMUInstrPass::handleMainFunction(MachineFunction &MF) {
    MachineFunction::iterator MFI= MF.begin();
    MachineBasicBlock::iterator FirstI=MFI->begin();
    FirstI++;
    MachineInstr &MI=*FirstI;
    DebugLoc DL=MI.getDebugLoc();

    BuildMI(*MFI,FirstI,DL,TII->get(X86::MOV64ri),X86::R15)
            .addGlobalAddress(MF.getMMI().getModule()->getGlobalVariable("pagetable"));
    BuildMI(*MFI,FirstI,DL,TII->get(X86::MOV64rr),X86::RBP).addReg(X86::RSP);
    BuildMI(*MFI,FirstI,DL,TII->get(X86::AND64ri32),X86::RBP).addReg(X86::RBP).addImm(0xfff);
    BuildMI(*MFI,FirstI,DL,TII->get(X86::ADD64ri32),X86::RBP).addReg(X86::RBP).addImm(0xFFF000);

    return true;
}

bool InsertVirtMMUInstrPass::handleControlFlowInstruction(MachineBasicBlock &MBB, MachineBasicBlock::iterator MBI) {
    MachineFunction &MF = *MBB.getParent();
    const DebugLoc &DL = MBI->getDebugLoc();

    if(MBI->getDesc().isCall()){
       if( MBI->getOperand(0).isGlobal()){
          MBI->print(errs());
       } else if(MBI->getOperand(0).isReg()){
           //MBI->print(errs());
       } else{
           errs()<<__func__<<" call operand0 is not global or reg\n";
           MBI->print(errs());
       }

    }

}
bool InsertVirtMMUInstrPass::insertDereferenceInstr(MachineBasicBlock &MBB, MachineBasicBlock::iterator MBI) {
     int opcode = MBI->getOpcode();
    if(isPOPorPush(opcode)){
        return false;
    }
    int baseIndex = findMemOprandFirstIndex(*MBI);
    if(baseIndex<0){
        return false;
    }
    MachineOperand &base = MBI->getOperand(baseIndex);
    if(base.isReg() && base.getReg()!=X86::RSP && base.getReg()!=X86::R15 &&base.getReg()!=X86::R15D){
        MBI->print(errs());
        //BuildMI(MBB,MI,MI.getDebugLoc(),TII->get(X86::FNOP));
        //lea (%eax) %r14
        MachineInstrBuilder MIB = BuildMI(MBB,*MBI,MBI->getDebugLoc(),TII->get(X86::LEA64r),X86::R14);
        for(int i=0;i<5;i++){
            MIB->addOperand(MBI->getOperand(baseIndex+i));
        }
        //RORX R14 R14>>12
        BuildMI(MBB,*MBI,MBI->getDebugLoc(),TII->get(X86::RORX64ri),X86::R14)
                .addReg(X86::R14)
                .addImm(12);

        BuildMI(MBB,*MBI,MBI->getDebugLoc(),TII->get(X86::ADD64rm),X86::R14).addReg(X86::R14)
                .addReg(X86::R15D).addImm(8).addReg(X86::R14D).addImm(0).addReg(0);
        //RORX R14 R14>>52
        BuildMI(MBB,*MBI,MBI->getDebugLoc(),TII->get(X86::RORX64ri),X86::R14)
                //.addReg(X86::RSP).addImm(1).addReg(0).addImm(-8).addReg(0)
                .addReg(X86::R14)
                .addImm(52);

        MBI->getOperand(baseIndex + 0).setReg(X86::R14);
        MBI->getOperand(baseIndex + 1).setImm(0);
        MBI->getOperand(baseIndex + 2).setReg(0);
        MBI->getOperand(baseIndex + 3).ChangeToImmediate(0);

        //MIB.getInstr()->print(errs());
        return true;
    }
    return false;
}

bool InsertVirtMMUInstrPass::handleInstruction(MachineBasicBlock &MBB,MachineBasicBlock::iterator MBI) {
    //replace RSP register with RBP
    bool isModified= false;
    unsigned test=0;
    if(isLEA(MBI->getOpcode())){

    }
    if(MBI->mayLoadOrStore()){
        isModified=insertDereferenceInstr(MBB,*MBI);
        test++;
    }
    if(hasControlFlow(*MBI)){
        isModified=handleControlFlowInstruction(MBB,MBI);
        test++;
    }
    if(test==2){
        errs()<<__func__<<":\n";
        MBI->print(errs());
        errs()<<"         load from mm\n";
    }
    return isModified;
}
bool InsertVirtMMUInstrPass::alignStack(MachineBasicBlock &MBB, MachineBasicBlock::iterator MBI) {
    MachineFunction &MF = *MBB.getParent();
    const DebugLoc &DL = MBI->getDebugLoc();
    X86MachineFunctionInfo *X86FI = MF.getInfo<X86MachineFunctionInfo>();

    MachineInstrBuilder MIB;
    const BasicBlock *LLVM_BB=MBB.getBasicBlock();
    const int stackSize = MF.getFrameInfo().getStackSize();
    int newStackSize = stackSize + 8 ;
    if(newStackSize > 4096){
        errs()<<"This function frame size is bigger than 4096.\n";
        return false;
    }
    //measure basic block
    //if  (RBP&0xfff >= stackSize + 8)
    //             goto continueMBB
    //MBB:
    //    MOV  RBP R14
    //    SAL  R14 4
    //    CMP  stacksize<<4, R14W
    //    JAE   ReserveMBB
    //AlignMBB
    //    SHR R14 16
    //    ADD (R15D + R14D*8 -8) R14D    [R14-1] + R14-1 +1
    //    SHL 12 R14
    //    LEA R14+8 RSP
    //    PUSH rbp
    //    AND rbp, (%r15-8)
    //    RBP = RSP-(stackSize+8) //we do not need to minus the RET index
    //    JMP ContinueMBB
    //ReserveMBB
    //    PUSH rbp
    //    RBP = RBP - (stackSize +16)
    //ContinueMBB:
    //    [rest of original instruction]
    MachineBasicBlock *AlignMBB = MF.CreateMachineBasicBlock(LLVM_BB);
    MachineBasicBlock *ReserveMBB = MF.CreateMachineBasicBlock(LLVM_BB);
    MachineBasicBlock *ContinueMBB = MF.CreateMachineBasicBlock(LLVM_BB);

    MachineFunction::iterator MBBIter = std::next(MBB.getIterator());
    MF.insert(MBBIter,AlignMBB);
    MF.insert(MBBIter,ReserveMBB);
    MF.insert(MBBIter,ContinueMBB);

    //Split MBB and move the tail portion down to ContinueMBB
    MachineBasicBlock::iterator BeforeMBBI = std::prev(MBI);
    ContinueMBB->splice(ContinueMBB->begin(),&MBB,MBI,MBB.end());
    ContinueMBB->transferSuccessorsAndUpdatePHIs(&MBB);

    //MBB
    BuildMI(&MBB,DL,TII->get(X86::MOV64rr),X86::R14).addReg(X86::RBP);
    BuildMI(&MBB,DL,TII->get(X86::SHL64ri), X86::R14).addReg(X86::R14).addImm(4);
    //in fact, we may need more stack size than the newStackSize to store
    // return address ,stack perserved params etc.s
    BuildMI(&MBB,DL,TII->get(X86::CMP16ri)).addReg(X86::R14W).addImm((newStackSize+8)<<4);
    BuildMI(&MBB,DL,TII->get(X86::JAE_1)).addMBB(ReserveMBB);

    //AlignMBB
    BuildMI(AlignMBB,DL,TII->get(X86::SHR64ri),X86::R14).addReg(X86::R14).addImm(16);
    BuildMI(AlignMBB,DL,TII->get(X86::ADD64rm),X86::R14).addReg(X86::R14)
            .addReg(X86::R15D).addImm(8).addReg(X86::R14D).addImm(-8).addReg(0);
    BuildMI(AlignMBB,DL,TII->get(X86::SHL64ri),X86::R14).addReg(X86::R14).addImm(12);
    BuildMI(AlignMBB,DL,TII->get(X86::LEA64r),X86::RSP)
            .addReg(X86::R14).addImm(0).addReg(0).addImm(0).addReg(0);
    BuildMI(AlignMBB,DL,TII->get(X86::PUSH64r)).addReg(X86::RBP);
    BuildMI(AlignMBB,DL,TII->get(X86::AND64rm),X86::RBP).addReg(X86::RBP)
            .addReg(X86::R15).addImm(0).addReg(0).addImm(0).addReg(0);
    BuildMI(AlignMBB,DL,TII->get(X86::SUB64ri32),X86::RBP).addReg(X86::RBP)
            .addImm(newStackSize);
    BuildMI(AlignMBB,DL,TII->get(X86::JAE_1)).addMBB(ContinueMBB);
    //ReserveMBB
    BuildMI(ReserveMBB,DL,TII->get(X86::PUSH64r)).addReg(X86::RBP);
    BuildMI(ReserveMBB,DL,TII->get(X86::SUB64ri32),X86::RBP).addReg(X86::RBP)
            .addImm(newStackSize + 8);
    //ContinueMBB
    MachineBasicBlock::iterator ContinueMBBI = ContinueMBB->getFirstNonPHI();

    //Add the control flow edges we need
    MBB.addSuccessor(ReserveMBB);
    MBB.addSuccessor(AlignMBB);
    AlignMBB->addSuccessor(ContinueMBB);
    ReserveMBB->addSuccessor(ContinueMBB);

    return true;

}
bool InsertVirtMMUInstrPass::handlePrologue(MachineFunction::iterator MFI) {
    MachineBasicBlock::iterator MBBI = MFI->begin(),MBE = MFI->end();

    if(!isSUB(MBBI->getOpcode()) && MBBI->getOperand(0).isReg() && MBBI->getOperand(0).getReg() == X86::RSP){
        errs()<<__func__<<"=== "<<MFI->getParent()->getName()<<" do not start with RSP-offset\n";
    }
    alignStack(*MFI,MBBI);

    return false;

}
bool InsertVirtMMUInstrPass::handleBlock(MachineFunction::iterator MFI) {
    bool ismodified=false;
    MachineBasicBlock::iterator MBBI = MFI->begin(),MBE = MFI->end(),MBBI_temp;

    for(;MBBI!=MBE;MBBI++){
       ismodified |= handleInstruction(*MFI,MBBI);
    }

    return ismodified;
}

bool InsertVirtMMUInstrPass::epalignStack(MachineBasicBlock &MBB, MachineBasicBlock::iterator MBI) {
    MachineFunction &MF=*MBB.getParent();
    const DebugLoc &DL = MBI->getDebugLoc();
    MachineBasicBlock::iterator MBBI = MBB.getFirstTerminator();

    const int stackSize = MF.getFrameInfo().getStackSize();
    int newStackSize = stackSize + 8;
    if (newStackSize >4086)return false;
    if(stackSize ==0)
        return false;

    MachineBasicBlock::iterator FirstCSPop = MBBI;
    // Skip the callee-saved pop instructions.
    int poptimes=0;
    while (MBBI != MBB.begin()) {
        MachineBasicBlock::iterator PI = std::prev(MBBI);
        unsigned Opc = PI->getOpcode();

        if (Opc != X86::DBG_VALUE && !PI->isTerminator()) {
            if ((Opc != X86::POP32r || !PI->getFlag(MachineInstr::FrameDestroy)) &&
            (Opc != X86::POP64r || !PI->getFlag(MachineInstr::FrameDestroy)))
            break;
            FirstCSPop = PI;
            poptimes++;
        }
        --MBBI;
    }
    if(FirstCSPop!=MBB.begin()){
        MBBI=std::prev(FirstCSPop);
    } else{
        MBBI=FirstCSPop;
    }
    //delete add operations of rsp
    if(isADD(MBBI->getOpcode())
       && MBBI->getOperand(0).isReg()
       && MBBI->getOperand(0).getReg()==X86::RSP){
        MBB.erase(MBBI);
    } else{
        MBBI->print(errs());
        errs()<<"==================\n";
    }
    for(int i=0;i<poptimes;i++){
        if(isPOP(FirstCSPop->getOpcode())){
            unsigned savedRegister=FirstCSPop->getOperand(0).getReg();
            BuildMI(MBB,FirstCSPop,FirstCSPop->getDebugLoc(),TII->get(X86::MOV64rm),savedRegister)
                    .addReg(X86::RSP).addImm(0).addReg(0).addImm(stackSize - (poptimes-i)*8).addReg(0);
            MBBI=FirstCSPop++;
            MBB.erase(MBBI);
        }
    }
    if(isRET(FirstCSPop->getOpcode())){
        MBB.erase(FirstCSPop);
    }

    // MBB:
    // MOV [RSP + stackSize] RBP
    // RORX RBP 12 R14
    // ADD (R15 + R14D*8) RSP
    // RORX R14 52 RSP
    // RORX (RSP-8) 12 R14
    // ADD  (R15 + R14D*8) R14
    // RORX R14 52 R14
    // jmp R14

    BuildMI(&MBB,DL,TII->get(X86::MOV64rm),X86::RBP)
            .addReg(X86::RSP).addImm(0).addReg(0).addImm(stackSize).addReg(0);
    BuildMI(&MBB,DL,TII->get(X86::RORX64ri),X86::R14)
            .addReg(X86::RBP) .addImm(12);
    BuildMI(&MBB,DL,TII->get(X86::ADD64rm),X86::R14).addReg(X86::R14)
            .addReg(X86::R15D).addImm(8).addReg(X86::R14D).addImm(0).addReg(0);
    BuildMI(&MBB,DL,TII->get(X86::RORX64ri),X86::RSP)
            .addReg(X86::R14)
            .addImm(52);
    BuildMI(&MBB,DL,TII->get(X86::RORX64mi),X86::R14)
            .addReg(X86::RSP).addImm(0).addReg(0).addImm(-8).addReg(0)
            .addImm(12);
    BuildMI(&MBB,DL,TII->get(X86::ADD64rm),X86::R14).addReg(X86::R14)
            .addReg(X86::R15D).addImm(8).addReg(X86::R14D).addImm(0).addReg(0);
    BuildMI(&MBB,DL,TII->get(X86::RORX64ri),X86::R14)
            .addReg(X86::R14)
            .addImm(52);
    BuildMI(&MBB,DL,TII->get(X86::JMP64r))
            .addReg(X86::R14);
    return true;
}
bool InsertVirtMMUInstrPass::handleEPrologue(MachineFunction::iterator MFI) {
    MachineBasicBlock::iterator MBBI_temp,MBBE = MFI->getFirstTerminator();
    MBBI_temp=MBBE;

    if(!isRET(MBBE->getOpcode())){
        errs()<<__func__<<"=== "<<MFI->getParent()->getName()<<" do not end with ret\n";
        return false;
    }
    //BuildMI(*MFI,*MBBE,MBBE->getDebugLoc(),TII->get(X86::ADD64ri32),X86::RBP).addReg(X86::RBP)
     //       .addImm(MFI->getParent()->getFrameInfo().getStackSize()+8);
    epalignStack(*MFI,MBBI_temp);

    return true;
}

bool InsertVirtMMUInstrPass::runOnMachineFunction(MachineFunction &MF){
    MachineFunction::iterator MFI, MFB=MF.begin(), MFE=MF.end();
    TII = MF.getSubtarget().getInstrInfo();
    TRI = MF.getSubtarget().getRegisterInfo();
    X86MachineFunctionInfo *X86FI = MF.getInfo<X86MachineFunctionInfo>();

    if(MF.getName()=="fillPagetable")
        return false;
    if(MF.getName()=="main"){
        handleMainFunction(MF);
        return true;
    }

    errs()<<"Function Name:  "<<MF.getName()<<" "<<MF.getFrameInfo().getStackSize()<<" "<<X86FI->getCalleeSavedFrameSize()<<"\n";
    for(MFI=MF.begin();MFI!=MFE;MFI++){
        handleBlock(MFI);
    }
    handlePrologue(MFB);
    handleEPrologue(--MFI);

//    errs()<<"--------------------------------------------------------\n";
//    errs()<<"Instructions Number:   "<<increatNum<<"   instructions\n\n";
    return false;
}

FunctionPass *llvm::insertVirtualMMUinstruction(){
    return new InsertVirtMMUInstrPass();
}
char InsertVirtMMUInstrPass::ID = 0;
static RegisterPass<InsertVirtMMUInstrPass> X("machinecount", "machine count pass");